package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io"
	"io/ioutil"

	//"log"
	"math"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

var StatusCodeMap map[int64]int64
var Total int
var Failed int
var Task chan int
var Conc chan int

func RequestByProxy(url_addr, proxy_addr, method string, body *bytes.Reader) (*http.Response, error) {
	var request *http.Request
	if method == "GET" {
		request, _ = http.NewRequest(method, url_addr, nil)
	} else {
		request, _ = http.NewRequest(method, url_addr, body)
		request.Header.Set("Content-Type", "application/x-www-form-urlencoded; param=value")
	}
	request.Header.Set("User-Agent", "xab")
	request.Header.Set("Accept-Encoding", "compress, gzip")
	request.Header.Set("Connection", "close")
	request.Close = true

	var pool *x509.CertPool
	if request.URL.Scheme == "https" {
		pool = x509.NewCertPool()
		caCertPath := capath
		caCrt, err1 := ioutil.ReadFile(caCertPath)
		if err1 != nil {
			fmt.Printf("%s [error] read cert file fail: %v\n", time.Now().Format("2006-01-02T15:04:05"), err1)
		}
		pool.AppendCertsFromPEM(caCrt)
		//pool.AppendCertsFromPEM([]byte(castr))
	}
	var client *http.Client
	var conn net.Conn
	var err error
	if proxy_addr != "" {
		proxy_addr = "http://" + proxy_addr
		proxy, err := url.Parse(proxy_addr)
		if err != nil {
			return nil, err
			fmt.Printf("%s ERR:GetByProxy err,ip:%s,err:\n", time.Now().Format("2006-01-02T15:04:05"), proxy_addr, err)
		}
		client = &http.Client{
			Transport: &http.Transport{
				Proxy: http.ProxyURL(proxy),
				Dial: func(netw, proxy string) (net.Conn, error) {
					deadline := time.Now().Add(15 * time.Second)
					conn, err = net.DialTimeout(netw, proxy, time.Second*15)
					if err != nil {
						return nil, err
					}
					conn.SetDeadline(deadline)
					remoteaddr = conn.RemoteAddr().String()
					return conn, nil
				},
				ResponseHeaderTimeout: time.Second * 15,
				TLSClientConfig:       &tls.Config{RootCAs: pool},
			},
		}
	} else {
		client = &http.Client{
			Transport: &http.Transport{
				Dial: func(netw, addr string) (net.Conn, error) {
					deadline := time.Now().Add(15 * time.Second)
					conn, err = net.DialTimeout(netw, addr, time.Second*15)
					if err != nil {
						return nil, err
					}
					remoteaddr = conn.RemoteAddr().String()
					conn.SetDeadline(deadline)
					return conn, nil
				},
				ResponseHeaderTimeout: time.Second * 15,
				TLSClientConfig:       &tls.Config{RootCAs: pool},
			},
		}
	}
	return client.Do(request)
}

func Get(url string) {
	Conc <- 1
	Total++
	startTime := time.Now().UnixNano()
	resp, err := RequestByProxy(url, proxyip, "GET", nil)
	endTime := time.Now().UnixNano()
	reqtime := float64(endTime-startTime) / 1000000000
	if err != nil {
		fmt.Printf("%s %s %.3f %s err:%v\n", time.Now().Format("2006-01-02T15:04:05"), remoteaddr, math.Trunc((reqtime+0.5/math.Pow10(3))*math.Pow10(3))/math.Pow10(3), url, err)
		Failed++
		Task <- 1
		<-Conc
		return
	}
	defer resp.Body.Close()
	io.Copy(ioutil.Discard, resp.Body)
	sc := int64(resp.StatusCode)
	fmt.Printf("%s %s %d %.3f %s\n", time.Now().Format("2006-01-02T15:04:05"), remoteaddr, sc, math.Trunc((reqtime+0.5/math.Pow10(3))*math.Pow10(3))/math.Pow10(3), url)
	Task <- 1
	<-Conc
}

func LoadUrls(fileName string, conc int) []string {
	result := make([]string, 0, 0)
	//_, err := url.Parse(fileName)
	//if err != nil {
	if !strings.HasPrefix(fileName, "http") {
		file, err1 := os.Open(fileName)
		if err1 != nil {
			fmt.Printf("%s [error] open file error,err:%s", time.Now().Format("2006-01-02T15:04:05"), err1)
			ShowHelp()
			os.Exit(1)
		}
		defer file.Close()
		bi := bufio.NewReader(file)
		for {
			line, err := bi.ReadString('\n')
			if err != nil {
				break
			}
			result = append(result, strings.TrimRight(line, "\r\n"))
		}
	} else {
		//result = []string{fileName}
		if conc >= 1 && xtotal == 0 {
			for i := 0; i < conc; i++ {
				result = append(result, fileName)
			}
		} else {
			result = []string{fileName}
		}
	}

	if xtotal > 0 {
		tmpresult := result
		for i := 1; i < xtotal; i++ {
			result = append(result, tmpresult...)
		}
	}
	return result
}

func ShowHelp() {
	fmt.Printf("Version: v1.1\n")
	fmt.Println("BuildTime: " + VERSION)
	fmt.Println("Usage: ./xab -c 10 urls.txt")
	fmt.Println("Options description:")
	fmt.Println("\t-c    Thread Num, default 1")
	fmt.Println("\t-n    Replay times")
	fmt.Println("\t-t    run how long time:seconds")
	fmt.Println("\t-x    set proxy")
	fmt.Println("\t-ca   ca.crt path,defautl:/etc/pki/tls/certs/ca-bundle.crt")
	fmt.Println("\t-h    Show help, default false")
	fmt.Println("\t-help Same as -h")
}

var proxyip, capath, remoteaddr string
var xtotal int
var VERSION string

func main() {
	help := flag.Bool("help", false, "Show help")
	h := flag.Bool("h", false, "Show help")
	c := flag.Int("c", 1, "Thread Num")
	totalnum := flag.Int("n", 0, "Replay times")
	times := flag.Int64("t", 0, "run how long time:seconds")
	xproxy := flag.String("x", "", "set proxy")
	ca := flag.String("ca", "/etc/pki/tls/certs/ca-bundle.crt", "ca.crt path")
	flag.Parse()

	if *help || *h {
		ShowHelp()
		return
	}

	xtotal = *totalnum
	proxyip = *xproxy
	capath = *ca

	urls := LoadUrls(flag.Arg(0), *c)

	StatusCodeMap = make(map[int64]int64)
	Task = make(chan int, *c)
	Conc = make(chan int, *c)
	timeend := time.Now().Unix() + *times

	//	var totalloop, ii int
	if *times > 0 {
		for {

			if timeend-time.Now().Unix() <= 0 {
				break
			}

			for i := 0; i < len(urls); i++ {
				go Get(urls[i])
			}

			for i := 0; i < len(urls); i++ {
				<-Task
			}

		}

	} else {

		for i := 0; i < len(urls); i++ {
			go Get(urls[i])
		}

		for i := 0; i < len(urls); i++ {
			<-Task
		}
	}
	fmt.Printf("%d of %d successed.\n", (Total - Failed), Total)

}
